import numpy as np
import matplotlib.pyplot as plt
from scipy.special import wofz

# --- Plotting Setup ---
# Use LaTeX for rendering text for a professional, publication-quality look
plt.rc('text', usetex=True)
plt.rc('font', family='serif', serif=['Computer Modern'])

# FIX: Add a preamble to load necessary LaTeX math packages (amsmath for \boldsymbol)
plt.rcParams['text.latex.preamble'] = r'''
\usepackage{amsmath}
\usepackage{amssymb}
\usepackage{bm}
'''

# Set font sizes for labels and titles for readability
plt.rc('xtick', labelsize=12)
plt.rc('ytick', labelsize=12)
plt.rc('axes', labelsize=14, titlesize=16)

# --- Physics Parameters ---
# Set constants for the plot, matching the paper's examples
a = 1.0
omega = 1.0

# Define the range of dimensionless interaction times (T*ω) to plot
T_omega_values = [0.5, 1, 2, 5, 10]

# Create the domain for the Unruh frequency Ω
Omega = np.linspace(-2.0, 4.0, 800)

# --- Core Calculation Function ---
def calculate_lineshape_squared(Omega, T, a, omega):
    """
    Calculates the squared probability amplitude along the gate.
    This functional form is identical for both the RR (Ω'=-Ω) and RL (Ω'=Ω) cuts.
    The result is proportional to: (Ω/sinh(πΩ))^2 * |w(T(aΩ-ω))|^2
    """
    # 1. Calculate the Unruh thermal prefactor: (Ω/sinh(πΩ))^2
    # This term is IR-safe; we handle the Ω=0 limit numerically to avoid division by zero.
    # The limit of (x/sinh(ax)) as x->0 is 1/a. Here, a=π.
    unruh_prefactor = np.zeros_like(Omega)
    nonzero_mask = Omega != 0
    unruh_prefactor[nonzero_mask] = (Omega[nonzero_mask] / np.sinh(np.pi * Omega[nonzero_mask]))**2
    unruh_prefactor[~nonzero_mask] = (1.0 / np.pi)**2 # Apply the known limit at Ω=0

    # 2. Define the argument for the Faddeeva function
    z = T * (a * Omega - omega)
    
    # 3. Calculate the Faddeeva function squared modulus: |w(z)|^2
    # scipy.special.wofz(z) computes w(z) = exp(-z^2)erfc(-iz)
    faddeeva_squared = np.abs(wofz(z))**2
    
    return unruh_prefactor * faddeeva_squared

# --- Plot Generation ---
fig, ax = plt.subplots(figsize=(10, 6))

# Loop through each T*ω value to plot a separate curve
for T_omega in T_omega_values:
    # Calculate the corresponding interaction duration T
    T = T_omega / omega
    
    # Calculate the probability density
    prob_density = calculate_lineshape_squared(Omega, T, a, omega)
    
    # Normalize each curve to its own maximum to emphasize the lineshape
    normalized_prob_density = prob_density / np.max(prob_density)
    
    # Plot the result
    ax.plot(Omega, normalized_prob_density, label=rf'$T\omega = {T_omega}$')

# --- Styling and Annotations ---
# Set the y-axis to a logarithmic scale to see the tails
ax.set_yscale('log')

# Add a vertical dashed line to mark the eternal-limit resonance
resonance_freq = omega / a
ax.axvline(x=resonance_freq, color='k', linestyle='--', linewidth=1.2, label=rf'Resonance $\Omega = \omega/a$')

# Set labels and title using LaTeX strings with bolding
ax.set_xlabel(r'\textbf{Unruh Frequency} $\boldsymbol{\Omega}$')
ax.set_ylabel(r'\textbf{Normalized Probability Density}')
ax.set_title(r'\textbf{1D Resonant Cuts Along Gates for RR and RL Channels}')

# Configure axes and legend
ax.set_xlim(Omega.min(), Omega.max())
ax.set_ylim(1e-4, 2.0)
ax.grid(True, which='both', linestyle='--', linewidth=0.5, alpha=0.7)
ax.legend(fontsize=12)

# Ensure tight layout to prevent labels from being cut off
plt.tight_layout()

# --- Save the Figure ---
# Save the plot as a high-resolution PNG file
output_filename = 'resonant_cuts_1D.png'
plt.savefig(output_filename, dpi=300, bbox_inches='tight')

print(f"✅ Plot saved successfully as '{output_filename}'")

# Display the plot
plt.show()